# 设置坡度和坡向范围

import os
import numpy as np
import pandas as pd
import rasterio
from rasterio.warp import reproject, Resampling
from sklearn.linear_model import LinearRegression
import matplotlib.pyplot as plt
from matplotlib import rcParams
rcParams['font.sans-serif'] = ['Noto Sans CJK JP']  # 选择简体中文
rcParams['axes.unicode_minus'] = False  # 解决负号显示问题
# ---------- 参数 ----------
dem_range = (3,3)
slope_range = (4,4)   # 坡度范围（单位：度）
aspect_range = (6,6)  # 坡向范围（单位：度）
high_years = [2010, 2012, 2013, 2014, 2015, 2018]
low_years = [2005, 2006, 2007, 2011, 2017, 2019, 2021, 2022]
mid_years = [y for y in range(2002, 2023) if y not in high_years + low_years]

# ---------- 函数 ----------
def read_raster(file_path):
    with rasterio.open(file_path) as src:
        return src.read(1).astype(np.float32), src.transform

def clean_raster_data(data, nodata_value):
    data[data == nodata_value] = np.nan
    # data[data<0] = np.nan
    return data

def resample_raster(src_data, src_transform, target_transform, target_shape):
    dst_data = np.empty(target_shape, dtype=np.float32)
    reproject(
        source=src_data,
        destination=dst_data,
        src_transform=src_transform,
        dst_transform=target_transform,
        src_crs="EPSG:4326",
        dst_crs="EPSG:4326",
        resampling=Resampling.nearest,
    )
    # dst_data[np.isnan(src_data)] = np.nan
    return dst_data

def get_coordinates(shape, transform):
    rows, cols = np.indices(shape)
    xs, ys = rasterio.transform.xy(transform, rows, cols)
    return np.array(xs), np.array(ys)

def extract_values_at_points(points, data, transform):
    rows, cols = rasterio.transform.rowcol(transform, points[:, 0], points[:, 1])
    valid_mask = (
        (rows >= 0) & (rows < data.shape[0]) &
        (cols >= 0) & (cols < data.shape[1])
    )
    values = np.full(len(points), np.nan)
    values[valid_mask] = data[rows[valid_mask], cols[valid_mask]]
    return values

def calculate_npp_by_year_group(years, base_shape, base_transform, dem_transform):
    npp_group_sum = np.zeros(base_shape, dtype=np.float32)
    count = np.zeros(base_shape, dtype=np.uint16)
    for year in years:
        data, npp_transform = read_raster(os.path.join('/home/gsr/Documents/代做外包/mean_npp',
                                f"MOD17A3HGF.061_Npp_500m_doy{year}001_aid0001.tif"))
        data = clean_raster_data(data, 32767)
        resampled = resample_raster(data, npp_transform, dem_transform,  base_shape)
        resampled /= 10
        mask_valid = ~np.isnan(resampled)
        npp_group_sum[mask_valid] += resampled[mask_valid]
        count[mask_valid] += 1
    avg = np.full_like(npp_group_sum, np.nan)
    avg[count > 0] = npp_group_sum[count > 0] / count[count > 0]
    return avg

from sklearn.metrics import mean_squared_error, mean_absolute_error

def run_regression(elevation, npp_vals, label):
    mask = ~np.isnan(elevation) & ~np.isnan(npp_vals)
    X = elevation[mask].reshape(-1, 1)
    y = npp_vals[mask]

    if len(y) < 10:
        print(f"[{label}] 样本过少（{len(y)}），跳过回归分析")
        return None

    model = LinearRegression().fit(X, y)
    y_pred = model.predict(X)

    errors = y_pred - y
    max_error = np.max(np.abs(errors))
    min_error = np.min(np.abs(errors))
    mean_error = np.mean(np.abs(errors))
    mse = mean_squared_error(y, y_pred)
    std_dev = np.std(errors)
    r2 = model.score(X, y)

    a = model.coef_[0]
    b = model.intercept_

    print(f"【{label}】回归分析结果：")
    print(f"  回归函数         : y = {a:.6f} * x + {b:.6f}")
    print(f"  R² 拟合优度       : {r2:.4f}")
    print(f"  样本数           : {len(y)}")
    print(f"  最大误差         : {max_error:.6f}")
    print(f"  最小误差         : {min_error:.6f}")
    print(f"  平均误差         : {mean_error:.6f}")
    print(f"  均方误差 (MSE)   : {mse:.6f}")
    print(f"  误差标准差       : {std_dev:.6f}")
    print("-" * 50)

    sample_size = min(1000, len(y))
    indices = np.random.choice(len(y), size=sample_size, replace=False)
    X_sample = X[indices]
    y_sample = y[indices]

    # 绘制散点图和拟合直线
    plt.figure(figsize=(8, 5))
    plt.scatter(X_sample, y_sample, s=10, alpha=0.5, label='样本点')
    plt.plot(X_sample, model.predict(X_sample), color='red', label='拟合直线')
    plt.xlabel('Elevation')
    plt.ylabel('NPP')
    plt.title(f'{label} 回归分析')
    plt.legend()
    plt.grid(True)
    plt.tight_layout()
    plt.show()

    return model


# ---------- 主流程 ----------
# dem_data, dem_transform = read_raster('/home/gsr/Documents/代做外包/SRTM地形/SRTMHB/srtm_hb.tif')
dem_data, dem_transform = read_raster('/home/gsr/Documents/代做外包/SRTM地形/SRTMDEM/srtm_dem.tif')
dem_data = clean_raster_data(dem_data, 32767)
# dem_data = clean_raster_data(dem_data, 65535)#-32768)

slope_data, _ = read_raster('/home/gsr/Documents/代做外包/SRTM地形/SRTMSLOPE/srtm_slope.tif')
# slope_data, _ = read_raster(os.path.join('/home/gsr/Documents/代做外包/SRTM地形/','Slope_csj.tif'))
# slope_data = clean_raster_data(slope_data, -3.4028231e+38)
slope_data[slope_data<0] = np.nan
slope_resampled = resample_raster(slope_data, _, dem_transform, dem_data.shape)

aspect_data, _ = read_raster('/home/gsr/Documents/代做外包/SRTM地形/SRTMSPECT/srtm_spect.tif')
aspect_data = clean_raster_data(aspect_data, 15)
aspect_resampled = resample_raster(aspect_data, _, dem_transform, dem_data.shape)

xs, ys = get_coordinates(dem_data.shape, dem_transform)
flat_data = {
    'x': xs.ravel(), 'y': ys.ravel(),
    'elevation': dem_data.ravel(),
    'slope': slope_resampled.ravel(),
    'aspect': aspect_resampled.ravel()
}
df = pd.DataFrame(flat_data)

# 按坡度和坡向筛选
condition = (
    (df['slope'] >= slope_range[0]) & (df['slope'] <= slope_range[1]) &
    (df['aspect'] >= aspect_range[0]) & (df['aspect'] <= aspect_range[1])
)
# condition = (
#     (df['elevation'] >= dem_range[0]) & (df['elevation'] <= dem_range[1]) &
#     (df['aspect'] >= aspect_range[0]) & (df['aspect'] <= aspect_range[1])
# )
df_filtered = df[condition]
points = df_filtered[['x', 'y']].values

# 平均NPP图像
npp_high = calculate_npp_by_year_group(high_years, dem_data.shape, dem_transform, dem_transform)
npp_mid = calculate_npp_by_year_group(mid_years, dem_data.shape, dem_transform, dem_transform)
npp_low = calculate_npp_by_year_group(low_years, dem_data.shape, dem_transform, dem_transform)

# 提取值
# elev_vals = df_filtered['slope'].values
elev_vals = df_filtered['elevation'].values
npp_high_vals = extract_values_at_points(points, npp_high, dem_transform)
npp_mid_vals = extract_values_at_points(points, npp_mid, dem_transform)
npp_low_vals = extract_values_at_points(points, npp_low, dem_transform)

# 回归分析
run_regression(elev_vals, npp_high_vals, "高NPP年")
run_regression(elev_vals, npp_mid_vals, "中NPP年")
run_regression(elev_vals, npp_low_vals, "低NPP年")

